#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import time
try:
    import gradio as gr
except ImportError:
    raise ImportError("gradio is required. Please install it with: pip install gradio")

from utils import xhs_imitative_writing_hot_list, xhs_imitative_writing
from config_manager import config_manager

def content_generation_with_progress(text, custom_prompt, progress=gr.Progress()):
    """
    内容生成功能：调用xhs_imitative_writing_hot_list函数，带进度显示
    """
    if not text.strip():
        yield "请输入内容", "", "", "请输入内容"
        return
    
    try:
        progress(0, desc="开始生成内容...")
        yield "", "", "", "开始生成内容..."
        time.sleep(15)
        
        progress(0.2, desc="正在调用AI模型生成内容...")
        yield "", "", "", "正在调用AI模型生成内容..."
        time.sleep(15)
        
        progress(0.4, desc="正在改写正文...")
        yield "", "", "", "正在改写正文..."
        time.sleep(13)
        
        progress(0.6, desc="去重中...")
        yield "", "", "", "去重中..."
        time.sleep(15)
        
        progress(0.8, desc="正在提取标签...")
        yield "", "", "", "正在提取标签..."
        time.sleep(5)
        
        # 调用实际函数
        title, content, tags = xhs_imitative_writing_hot_list(text, custom_prompt)
        
        # 将标签列表转换为字符串显示
        tags_str = ", ".join(tags) if tags else ""
        
        progress(1.0, desc="生成完成！")
        yield title, content, tags_str, "生成完成！"
        
    except Exception as e:
        error_msg = f"生成失败: {str(e)}"
        yield error_msg, "", "", error_msg

def imitative_writing_with_progress(title, content, tags, custom_prompt, progress=gr.Progress()):
    """
    仿写功能：调用xhs_imitative_writing函数，带进度显示
    """
    if not title.strip() or not content.strip():
        yield "请输入标题和内容", "", "", "请输入标题和内容"
        return
    
    try:
        progress(0, desc="开始仿写...")
        yield "", "", "", "开始仿写..."
        time.sleep(10)
        
        progress(0.3, desc="正在调用AI模型进行仿写...")
        yield "", "", "", "正在调用AI模型进行仿写..."
        time.sleep(15)
        
        progress(0.6, desc="正在违规检测...")
        yield "", "", "", "正在违规检测..."
        time.sleep(15)
        
        progress(0.9, desc="正在提取标签...")
        yield "", "", "", "正在提取标签..."
        time.sleep(15)
        
        # 将标签字符串转换为列表
        tags_list = [tag.strip() for tag in tags.split(",") if tag.strip()] if tags else []
        
        # 调用实际函数
        new_title, new_content, new_tags = xhs_imitative_writing(title, content, tags_list, custom_prompt)
        
        # 将标签列表转换为字符串显示
        new_tags_str = ", ".join(new_tags) if new_tags else ""
        
        progress(1.0, desc="仿写完成！")
        yield new_title, new_content, new_tags_str, "仿写完成！"
        
    except Exception as e:
        error_msg = f"仿写失败: {str(e)}"
        yield error_msg, "", "", error_msg

# 创建Gradio界面
with gr.Blocks(title="小红书内容生成工具", theme=gr.themes.Soft()) as app:
    gr.Markdown("# 小红书内容生成工具")
    gr.Markdown("这个工具可以帮你生成和仿写小红书爆款文案")
    
    # 初始化配置（这些值会在页面加载时被更新）
    current_content_prompt, current_imitative_prompt = config_manager.load_user_config()
    
    # 使用State来保存当前的提示词配置
    content_prompt_state = gr.State(value=current_content_prompt)
    imitative_prompt_state = gr.State(value=current_imitative_prompt)
    
    with gr.Tabs():
        # 内容生成Tab
        with gr.TabItem("内容生成"):
            gr.Markdown("### 输入原始内容，生成小红书爆款文案")
            
            # 提示词编辑区域
            with gr.Accordion("提示词设置", open=False):
                content_generation_prompt = gr.Textbox(
                    label="内容生成提示词",
                    value=current_content_prompt,
                    lines=15,
                    max_lines=25,
                    placeholder="在这里编辑内容生成的提示词..."
                )
                with gr.Row():
                    reset_content_btn = gr.Button("重置为默认提示词", variant="secondary")
                    save_content_btn = gr.Button("保存提示词", variant="primary")
                
                # 保存状态显示
                save_status_content = gr.Textbox(
                    label="保存状态",
                    value="",
                    interactive=False,
                    lines=1,
                    visible=False
                )
            
            with gr.Row():
                with gr.Column(scale=2):
                    input_text = gr.Textbox(
                        label="输入内容",
                        placeholder="请输入要生成文案的原始内容...",
                        lines=8,
                        max_lines=15
                    )
                    # 添加状态显示
                    status_display = gr.Textbox(
                        label="处理状态",
                        value="等待开始...",
                        interactive=False,
                        lines=1
                    )
                    generate_btn = gr.Button("生成", variant="primary", size="lg")
                
                with gr.Column(scale=3):
                    output_title = gr.Textbox(
                        label="生成的标题",
                        interactive=False,
                        lines=2
                    )
                    output_content = gr.Textbox(
                        label="生成的正文",
                        interactive=False,
                        lines=10,
                        max_lines=20
                    )
                    output_tags = gr.Textbox(
                        label="生成的标签",
                        interactive=False,
                        lines=2
                    )
        
        # 仿写Tab
        with gr.TabItem("仿写"):
            gr.Markdown("### 输入标题、内容和标签，进行仿写优化")
            
            # 提示词编辑区域
            with gr.Accordion("提示词设置", open=False):
                imitative_writing_prompt = gr.Textbox(
                    label="仿写提示词",
                    value=current_imitative_prompt,
                    lines=15,
                    max_lines=25,
                    placeholder="在这里编辑仿写的提示词..."
                )
                with gr.Row():
                    reset_imitative_btn = gr.Button("重置为默认提示词", variant="secondary")
                    save_imitative_btn = gr.Button("保存提示词", variant="primary")
                
                # 保存状态显示
                save_status_imitative = gr.Textbox(
                    label="保存状态",
                    value="",
                    interactive=False,
                    lines=1,
                    visible=False
                )
            
            with gr.Row():
                with gr.Column(scale=2):
                    rewrite_title = gr.Textbox(
                        label="标题",
                        placeholder="请输入标题...",
                        lines=2
                    )
                    rewrite_content = gr.Textbox(
                        label="内容",
                        placeholder="请输入正文内容...",
                        lines=8,
                        max_lines=15
                    )
                    rewrite_tags = gr.Textbox(
                        label="标签",
                        placeholder="请输入标签，用逗号分隔...",
                        lines=2
                    )
                    # 添加状态显示
                    rewrite_status_display = gr.Textbox(
                        label="处理状态",
                        value="等待开始...",
                        interactive=False,
                        lines=1
                    )
                    rewrite_btn = gr.Button("仿写", variant="primary", size="lg")
                
                with gr.Column(scale=3):
                    rewrite_output_title = gr.Textbox(
                        label="仿写后的标题",
                        interactive=False,
                        lines=2
                    )
                    rewrite_output_content = gr.Textbox(
                        label="仿写后的正文",
                        interactive=False,
                        lines=10,
                        max_lines=20
                    )
                    rewrite_output_tags = gr.Textbox(
                        label="仿写后的标签",
                        interactive=False,
                        lines=2
                    )
    
    # 事件处理函数
    def reset_content_generation_prompt():
        """重置内容生成提示词为默认值"""
        default_prompt = config_manager.get_default_content_generation_prompt()
        return default_prompt, default_prompt
    
    def reset_imitative_writing_prompt():
        """重置仿写提示词为默认值"""
        default_prompt = config_manager.get_default_imitative_writing_prompt()
        return default_prompt, default_prompt
    
    def save_content_prompt_handler(prompt, current_imitative):
        """保存内容生成提示词"""
        result = config_manager.save_user_config(prompt, current_imitative)
        return result, gr.update(visible=True), prompt
    
    def save_imitative_prompt_handler(prompt, current_content):
        """保存仿写提示词"""
        result = config_manager.save_user_config(current_content, prompt)
        return result, gr.update(visible=True), prompt
    
    def update_content_prompt_state(prompt):
        """更新内容生成提示词状态"""
        return prompt
    
    def update_imitative_prompt_state(prompt):
        """更新仿写提示词状态"""
        return prompt
    
    # 绑定事件
    
    # 页面加载事件 - 确保每次页面刷新时都重新加载配置
    def load_page_config():
        """页面加载时重新读取配置文件"""
        content_prompt, imitative_prompt = config_manager.load_user_config()
        return content_prompt, imitative_prompt, content_prompt, imitative_prompt
    
    # 页面加载时更新所有相关组件
    app.load(
        load_page_config,
        outputs=[content_generation_prompt, imitative_writing_prompt, content_prompt_state, imitative_prompt_state]
    )
    
    # 事件处理
    reset_content_btn.click(
        reset_content_generation_prompt,
        outputs=[content_generation_prompt, content_prompt_state]
    )
    
    reset_imitative_btn.click(
        reset_imitative_writing_prompt,
        outputs=[imitative_writing_prompt, imitative_prompt_state]
    )
    
    # 保存按钮事件
    save_content_btn.click(
        save_content_prompt_handler,
        inputs=[content_generation_prompt, imitative_prompt_state],
        outputs=[save_status_content, save_status_content, content_prompt_state]
    )
    
    save_imitative_btn.click(
        save_imitative_prompt_handler,
        inputs=[imitative_writing_prompt, content_prompt_state],
        outputs=[save_status_imitative, save_status_imitative, imitative_prompt_state]
    )
    
    # 提示词变化时更新状态
    content_generation_prompt.change(
        update_content_prompt_state,
        inputs=[content_generation_prompt],
        outputs=[content_prompt_state]
    )
    
    imitative_writing_prompt.change(
        update_imitative_prompt_state,
        inputs=[imitative_writing_prompt],
        outputs=[imitative_prompt_state]
    )
    
    # 生成和仿写按钮事件
    generate_btn.click(
        content_generation_with_progress,
        inputs=[input_text, content_prompt_state],
        outputs=[output_title, output_content, output_tags, status_display]
    )
    
    rewrite_btn.click(
        imitative_writing_with_progress,
        inputs=[rewrite_title, rewrite_content, rewrite_tags, imitative_prompt_state],
        outputs=[rewrite_output_title, rewrite_output_content, rewrite_output_tags, rewrite_status_display]
    )

if __name__ == "__main__":
    app.launch(
        server_name="0.0.0.0",
        server_port=8032,
        share=False,
        debug=True
    )